import torch
import torch.nn as nn
import torch.distributions as D


class NormalSample(nn.Module):
    """正态分布采样类
    Args:
        mean ([Tensor]): [正态分布的均值]
        std ([Tensor]): [正态分布的标准差]
        observation ([Tensor]): [分布的观测, 可事先指定观测, 观测可有可无]
    """
    def __init__(self,
                 mean,
                 std,
                 observation,
                 ):
        super().__init__()
        self.mean = mean
        self.std = std
        self.register_buffer("epslion", torch.randn(self.mean.size()))
        self.observation = observation
        self.samples = None

    def sample(self):
        """参数采样
        Returns:
            权重参数张量: 每个参数均服从正态分布，从该分布中采样
        """
        self.epslion.data.normal_()
        self.samples = self.mean + self.epslion * self.std

    def log_prob(self):
        """计算对数似然 log(w|\theta) or log(w)
        Returns:
            Tensor: 所有参数的对数似然的总和
        """
        _normal = D.Normal(self.mean.data, self.std.data)
        if self.observation is not None:
            _log_prob = _normal.log_prob(self.observation).sum()
        else:
            _log_prob = _normal.log_prob(self.samples).sum()
        return _log_prob

    @property
    def get_samples(self):
        return self.samples

    def set_observation(self, observation):
        self.observation = observation
